#include "RsaCrypto.h"
#include "openssl/rsa.h"
#include "openssl/pem.h"
#include "openssl/err.h"
#include <assert.h>
#include "Base64.h"

// 构造对象并加载秘钥文件中的数据到内存中
RsaCrypto::RsaCrypto(string fileName, KeyType type) {
    // 以二进制读形式打开秘钥文件
    BIO* bio = BIO_new_file(fileName.data(), "rb");
    assert(bio);

    // 根据type确定读取公钥还是私钥
    if(type == PublicKey) { PEM_read_bio_PUBKEY(bio, &_pubKey, NULL, NULL); }
    else { PEM_read_bio_PrivateKey(bio, &_priKey, NULL, NULL); }
    BIO_free(bio);
}

RsaCrypto::~RsaCrypto() {
    if(_pubKey) { EVP_PKEY_free(_pubKey); }
    else { EVP_PKEY_free(_priKey); }
}

// 将秘钥字符串data解析成秘钥类型
void RsaCrypto::parseStringToKey(string data, KeyType type) {
    BIO* bio = BIO_new_mem_buf(data.data(), data.size());
    assert(bio);

    // 根据type确定读取公钥还是私钥
    if(type == PublicKey) { PEM_read_bio_PUBKEY(bio, &_pubKey, NULL, NULL); }
    else { PEM_read_bio_PrivateKey(bio, &_priKey, NULL, NULL); }
    BIO_free(bio);
}

// 生成密钥对
void RsaCrypto::generateRsaKey(KeyLength bits, string pub, string pri) {
    // 实例化[密钥对上下文]: 采用RSA非对称加密算法
    EVP_PKEY_CTX* ctx = EVP_PKEY_CTX_new_id(EVP_PKEY_RSA, NULL);

    // 设置[密钥对的生成操作]的参数和上下文
    int ret = EVP_PKEY_keygen_init(ctx);
    assert(ret == 1);

    // 指定密钥对的长度
    ret = EVP_PKEY_CTX_set_rsa_keygen_bits(ctx, bits);
    assert(ret == 1);

    // 生成密钥对后, 释放密钥对上下文: 用_priKey存储生成的密钥对
    ret = EVP_PKEY_generate(ctx, &_priKey);
    assert(ret == 1);
    EVP_PKEY_CTX_free(ctx);

    // 将私钥写入磁盘文件
    BIO* bio = BIO_new_file(pri.data(), "wb"); // 以二进制写入方式打开私钥文件
    ret = PEM_write_bio_PrivateKey(bio, _priKey, NULL, NULL, 0, NULL, NULL); // 写入时, 不加密私钥, 也不设置执行加密所需的密码
    assert(ret == 1);
    BIO_flush(bio); // 刷新BIO对象: 将私钥立即从内存中刷新到磁盘文件中
    BIO_free(bio);

    // 将公钥写入磁盘文件
    bio = BIO_new_file(pub.data(), "wb");        // 以二进制写入方式打开私钥文件
    ret = PEM_write_bio_PUBKEY(bio, _priKey); // 写入时, 不加密私钥, 也不设置执行加密所需的密码
    assert(ret == 1);
    BIO_flush(bio); // 刷新BIO对象: 将私钥立即从内存中刷新到磁盘文件中
    BIO_free(bio);
}

// 用公钥加密[明文data]得到密文, 并返回[Base64对密文进行编码之后的结果]
string RsaCrypto::pubKeyEncrypt(string data) {
    // 创建与[给定密钥对象]相关联的上下文
    assert(_pubKey);
    EVP_PKEY_CTX* ctx = EVP_PKEY_CTX_new(_pubKey, NULL);
    assert(ctx);

    // 初始化[加密操作的上下文], 并设置填充模式
    int ret = EVP_PKEY_encrypt_init(ctx);
    assert(ret == 1);
    ret = EVP_PKEY_CTX_set_rsa_padding(ctx, RSA_PKCS1_OAEP_PADDING);
    assert(ret == 1);

    // 使用公钥加密data
    size_t outLen = 0;
    ret = EVP_PKEY_encrypt(ctx, NULL, &outLen, reinterpret_cast<const unsigned char*>(data.data()), data.size()); // 获取outLen
    assert(ret == 1);
    unsigned char* out = new unsigned char[outLen];
    ret = EVP_PKEY_encrypt(ctx, out, &outLen, reinterpret_cast<const unsigned char*>(data.data()), data.size());
    assert(ret == 1);

    // 用Base64对密文进行编码
    Base64 base;
    string retStr = base.encode(reinterpret_cast<char*>(out), outLen);

    // 释放资源并返回密文
    delete[] out;
    EVP_PKEY_CTX_free(ctx);
    return retStr;
}

// 用私钥解密[密文data], 返回明文
string RsaCrypto::priKeyDecrypt(string data) {
    // 创建与[给定密钥对象]相关联的上下文
    assert(_priKey);
    EVP_PKEY_CTX* ctx = EVP_PKEY_CTX_new(_priKey, NULL);
    assert(ctx);

    // 初始化[解密操作的上下文], 并设置填充模式
    int ret = EVP_PKEY_decrypt_init(ctx);
    assert(ret == 1);
    ret = EVP_PKEY_CTX_set_rsa_padding(ctx, RSA_PKCS1_OAEP_PADDING);
    assert(ret == 1);

    // 用Base64解码data
    Base64 base;
    data = base.decode(data);

    // 使用私钥解密data
    size_t outLen = 0;
    ret = EVP_PKEY_decrypt(ctx, NULL, &outLen, reinterpret_cast<const unsigned char*>(data.data()), data.size()); // 获取outLen
    assert(ret == 1);

    unsigned char* out = new unsigned char[outLen];
    ret = EVP_PKEY_decrypt(ctx, out, &outLen, reinterpret_cast<const unsigned char*>(data.data()), data.size());
    assert(ret == 1);

    // 释放资源并返回明文
    string retStr(reinterpret_cast<char*>(out), outLen);
    delete[] out;
    EVP_PKEY_CTX_free(ctx);
    return retStr;
}

// 数据签名: hashMethod用于指定要使用的哈希算法
string RsaCrypto::sign(string data, HashType hashMethod) {
    // 计算data的哈希值
    Hash hash(hashMethod);
    hash.addData(data);
    string md= hash.result(Hash::Type::Binary); // result()返回的哈希字符串是二进制格式的

    // 创建与[给定密钥对象]相关联的上下文
    assert(_priKey);
    EVP_PKEY_CTX* ctx = EVP_PKEY_CTX_new(_priKey, NULL);
    assert(ctx);

    // 初始化[签名操作的上下文], 并设置填充模式(签名和校验不能使用OAEP的填充方式)
    int ret = EVP_PKEY_sign_init(ctx);
    assert(ret == 1);
    ret = EVP_PKEY_CTX_set_rsa_padding(ctx, RSA_PKCS1_PADDING);
    assert(ret == 1);

    // 设置签名操作使用的哈希算法
    ret = EVP_PKEY_CTX_set_signature_md(ctx, HashMethods.at(hashMethod)());
    assert(ret == 1);

    // 使用私钥对[哈希值md]进行签名
    size_t outLen = 0;
    ret = EVP_PKEY_sign(ctx, NULL, &outLen, reinterpret_cast<const unsigned char*>(md.data()), md.size()); // 获取outLen
    assert(ret == 1);
    unsigned char* out = new unsigned char[outLen];
    ret = EVP_PKEY_sign(ctx, out, &outLen, reinterpret_cast<const unsigned char*>(md.data()), md.size());
    ERR_print_errors_fp(stderr);
    assert(ret == 1);

    // 用Base64对签名进行编码
    Base64 base;
    string retStr = base.encode(reinterpret_cast<char*>(out), outLen);

    // 释放资源并返回密文s
    delete[] out;
    EVP_PKEY_CTX_free(ctx);
    return retStr;
}

// 签名校验: sign为签名, data为待检验的数据 (sign是经过Base64编码的)
bool RsaCrypto::verify(string sign, string data, HashType hashMethod) {
    // 用Base64解码data
    Base64 base;
    sign = base.decode(sign);

    // 计算data的哈希值
    Hash hash(hashMethod);
    hash.addData(data);
    string md= hash.result(Hash::Type::Binary); // result()返回的哈希字符串是二进制格式的

    // 创建与[给定密钥对象]相关联的上下文
    assert(_pubKey);
    EVP_PKEY_CTX* ctx = EVP_PKEY_CTX_new(_pubKey, NULL);
    assert(ctx);

    // 初始化[校验签名操作的上下文], 并设置填充模式(签名和校验不能使用OAEP的填充方式)
    int ret = EVP_PKEY_verify_init(ctx);
    assert(ret == 1);
    ret = EVP_PKEY_CTX_set_rsa_padding(ctx, RSA_PKCS1_PADDING);
    assert(ret == 1);

    // 设置校验签名使用的哈希算法
    ret = EVP_PKEY_CTX_set_signature_md(ctx, HashMethods.at(hashMethod)());
    assert(ret == 1);

    // 签名校验
    size_t outLen = 0;
    ret = EVP_PKEY_verify(ctx, reinterpret_cast<const unsigned char*>(sign.data()), sign.size(),
                                reinterpret_cast<const unsigned char*>(md.data()), md.size()); // 获取outLen
    EVP_PKEY_CTX_free(ctx);
    if(ret == 1) { return true; }
    return false;
}
